#!/bin/bash
"""
DeepSpeed MOE + Qwen3-Coder API测试脚本
用于测试API的各项功能
"""

set -e

# 配置
API_BASE_URL="http://localhost:8000"
API_KEY="roo-code-key-2024"
MODEL_NAME="qwen3-coder-480b-a35b-instruct"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
log_step() { echo -e "${BLUE}[TEST]${NC} $1"; }

# 测试函数
test_health() {
    log_step "测试健康检查..."
    
    response=$(curl -s -w "\n%{http_code}" "$API_BASE_URL/health")
    status_code=$(echo "$response" | tail -n1)
    body=$(echo "$response" | head -n-1)
    
    if [ "$status_code" = "200" ]; then
        log_info "✓ 健康检查通过"
        echo "$body" | jq '.'
    else
        log_error "✗ 健康检查失败 (HTTP $status_code)"
        echo "$body"
        return 1
    fi
}

test_models() {
    log_step "测试模型列表..."
    
    response=$(curl -s -w "\n%{http_code}" \
        -H "Authorization: Bearer $API_KEY" \
        "$API_BASE_URL/v1/models")
    
    status_code=$(echo "$response" | tail -n1)
    body=$(echo "$response" | head -n-1)
    
    if [ "$status_code" = "200" ]; then
        log_info "✓ 模型列表获取成功"
        echo "$body" | jq '.'
    else
        log_error "✗ 模型列表获取失败 (HTTP $status_code)"
        echo "$body"
        return 1
    fi
}

test_simple_chat() {
    log_step "测试简单聊天功能..."
    
    local data='{
        "model": "'$MODEL_NAME'",
        "messages": [
            {"role": "user", "content": "Write a simple Python function to calculate factorial"}
        ],
        "max_tokens": 200,
        "temperature": 0.1
    }'
    
    response=$(curl -s -w "\n%{http_code}" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "$data" \
        "$API_BASE_URL/v1/chat/completions")
    
    status_code=$(echo "$response" | tail -n1)
    body=$(echo "$response" | head -n-1)
    
    if [ "$status_code" = "200" ]; then
        log_info "✓ 简单聊天测试通过"
        echo "$body" | jq '.choices[0].message.content'
        echo "$body" | jq '.usage'
    else
        log_error "✗ 简单聊天测试失败 (HTTP $status_code)"
        echo "$body"
        return 1
    fi
}

test_code_generation() {
    log_step "测试代码生成功能..."
    
    local data='{
        "model": "'$MODEL_NAME'",
        "messages": [
            {"role": "user", "content": "Create a Python class for a binary search tree with insert and search methods"}
        ],
        "max_tokens": 500,
        "temperature": 0.2
    }'
    
    response=$(curl -s -w "\n%{http_code}" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "$data" \
        "$API_BASE_URL/v1/chat/completions")
    
    status_code=$(echo "$response" | tail -n1)
    body=$(echo "$response" | head -n-1)
    
    if [ "$status_code" = "200" ]; then
        log_info "✓ 代码生成测试通过"
        echo "$body" | jq '.choices[0].message.content'
    else
        log_error "✗ 代码生成测试失败 (HTTP $status_code)"
        echo "$body"
        return 1
    fi
}

test_system_prompt() {
    log_step "测试系统提示功能..."
    
    local data='{
        "model": "'$MODEL_NAME'",
        "system": "You are a senior software engineer specializing in code review. Always provide detailed feedback.",
        "messages": [
            {"role": "user", "content": "Review this Python code: def divide(a, b): return a/b"}
        ],
        "max_tokens": 300
    }'
    
    response=$(curl -s -w "\n%{http_code}" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "$data" \
        "$API_BASE_URL/v1/chat/completions")
    
    status_code=$(echo "$response" | tail -n1)
    body=$(echo "$response" | head -n-1)
    
    if [ "$status_code" = "200" ]; then
        log_info "✓ 系统提示测试通过"
        echo "$body" | jq '.choices[0].message.content'
    else
        log_error "✗ 系统提示测试失败 (HTTP $status_code)"
        echo "$body"
        return 1
    fi
}

test_streaming() {
    log_step "测试流式输出功能..."
    
    local data='{
        "model": "'$MODEL_NAME'",
        "messages": [
            {"role": "user", "content": "Explain how a binary search algorithm works step by step"}
        ],
        "stream": true
    }'
    
    log_info "正在接收流式响应..."
    local start_time=$(date +%s)
    
    response=$(curl -s -N \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "$data" \
        "$API_BASE_URL/v1/chat/completions")
    
    local end_time=$(date +%s)
    local duration=$((end_time - start_time))
    
    if [ ! -z "$response" ]; then
        log_info "✓ 流式输出测试通过 (耗时: ${duration}秒)"
        echo "$response" | head -10
    else
        log_error "✗ 流式输出测试失败"
        return 1
    fi
}

test_authentication() {
    log_step "测试认证功能..."
    
    # 测试错误API Key
    local wrong_key="wrong-api-key"
    
    response=$(curl -s -w "\n%{http_code}" \
        -H "Authorization: Bearer $wrong_key" \
        "$API_BASE_URL/v1/models")
    
    status_code=$(echo "$response" | tail -n1)
    
    if [ "$status_code" = "401" ]; then
        log_info "✓ 认证测试通过 (正确拒绝了无效API Key)"
    else
        log_error "✗ 认证测试失败 - 应该返回401但返回了$status_code"
        return 1
    fi
    
    # 测试无认证
    response=$(curl -s -w "\n%{http_code}" \
        "$API_BASE_URL/v1/models")
    
    status_code=$(echo "$response" | tail -n1)
    
    if [ "$status_code" = "401" ]; then
        log_info "✓ 无认证测试通过"
    else
        log_error "✗ 无认证测试失败 - 应该返回401但返回了$status_code"
        return 1
    fi
}

test_error_handling() {
    log_step "测试错误处理..."
    
    # 测试无效模型
    local data='{
        "model": "invalid-model",
        "messages": [{"role": "user", "content": "Hello"}]
    }'
    
    response=$(curl -s -w "\n%{http_code}" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "$data" \
        "$API_BASE_URL/v1/chat/completions")
    
    status_code=$(echo "$response" | tail -n1)
    body=$(echo "$response" | head -n-1)
    
    if [ "$status_code" = "400" ]; then
        log_info "✓ 无效模型错误处理测试通过"
        echo "$body" | jq '.'
    else
        log_error "✗ 无效模型错误处理测试失败 (HTTP $status_code)"
        return 1
    fi
    
    # 测试空消息
    local empty_data='{
        "model": "'$MODEL_NAME'",
        "messages": []
    }'
    
    response=$(curl -s -w "\n%{http_code}" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "$empty_data" \
        "$API_BASE_URL/v1/chat/completions")
    
    status_code=$(echo "$response" | tail -n1)
    
    if [ "$status_code" = "400" ]; then
        log_info "✓ 空消息错误处理测试通过"
    else
        log_error "✗ 空消息错误处理测试失败"
        return 1
    fi
}

test_performance() {
    log_step "测试性能..."
    
    # 发送多个并发请求
    local concurrent_requests=3
    
    log_info "发送$concurrent_requests个并发请求..."
    
    for i in $(seq 1 $concurrent_requests); do
        (
            local data='{
                "model": "'$MODEL_NAME'",
                "messages": [{"role": "user", "content": "Write fibonacci function"}],
                "max_tokens": 100
            }'
            
            local start_time=$(date +%s%3N)
            response=$(curl -s -w "\n%{http_code}" \
                -H "Content-Type: application/json" \
                -H "Authorization: Bearer $API_KEY" \
                -d "$data" \
                "$API_BASE_URL/v1/chat/completions")
            local end_time=$(date +%s%3N)
            local duration=$((end_time - start_time))
            
            status_code=$(echo "$response" | tail -n1)
            
            if [ "$status_code" = "200" ]; then
                echo "Request $i: ✓ 成功 (${duration}ms)"
            else
                echo "Request $i: ✗ 失败 (${status_code})"
            fi
        ) &
    done
    
    wait
    log_info "✓ 性能测试完成"
}

# 主函数
main() {
    echo "==============================================="
    echo "   DeepSpeed MOE + Qwen3-Coder API测试套件"
    echo "==============================================="
    echo
    
    # 检查API服务器是否运行
    if ! curl -s "$API_BASE_URL/health" > /dev/null; then
        log_error "API服务器未运行，请先启动服务："
        log_error "  ./start_api.sh"
        exit 1
    fi
    
    log_info "开始API测试..."
    echo
    
    # 运行所有测试
    local tests_passed=0
    local tests_total=0
    
    # 测试列表
    local test_functions=(
        "test_health"
        "test_models" 
        "test_simple_chat"
        "test_code_generation"
        "test_system_prompt"
        "test_streaming"
        "test_authentication"
        "test_error_handling"
        "test_performance"
    )
    
    for test_func in "${test_functions[@]}"; do
        tests_total=$((tests_total + 1))
        echo "============================================="
        if $test_func; then
            tests_passed=$((tests_passed + 1))
            echo "测试 $tests_total/${#test_functions[@]} 通过"
        else
            echo "测试 $tests_total/${#test_functions[@]} 失败"
        fi
        echo
    done
    
    # 显示结果
    echo "==============================================="
    echo "           测试结果总结"
    echo "==============================================="
    echo "总测试数: $tests_total"
    echo "通过测试: $tests_passed"
    echo "失败测试: $((tests_total - tests_passed))"
    
    if [ $tests_passed -eq $tests_total ]; then
        log_info "🎉 所有测试通过！API系统运行正常。"
        exit 0
    else
        log_error "❌ 部分测试失败，请检查配置和日志。"
        exit 1
    fi
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        --url)
            API_BASE_URL="$2"
            shift 2
            ;;
        --key)
            API_KEY="$2"
            shift 2
            ;;
        --model)
            MODEL_NAME="$2"
            shift 2
            ;;
        --help|-h)
            echo "用法: $0 [选项]"
            echo "选项:"
            echo "  --url URL        API基础URL (默认: http://localhost:8000)"
            echo "  --key KEY        API密钥 (默认: roo-code-key-2024)"
            echo "  --model MODEL    模型名称 (默认: qwen3-coder-480b-a35b-instruct)"
            echo "  --help           显示帮助信息"
            exit 0
            ;;
        *)
            log_error "未知参数: $1"
            exit 1
            ;;
    esac
done

# 执行主函数
main "$@"